package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

type ACNode struct {
	children map[rune]*ACNode
	isEnd    bool
	//code     string
	fail *ACNode
	word string
}

func NewACNode( /*code string*/ ) *ACNode {
	return &ACNode{
		children: make(map[rune]*ACNode),
		isEnd:    false,
		word:     "",
		//code:     code,
	}
}

type ACAutomaton struct {
	root *ACNode
}

func NewACAutomaton() *ACAutomaton {
	//return &ACAutomaton{root: NewACNode("")}
	return &ACAutomaton{root: NewACNode()}
}

func (ac *ACAutomaton) AddWord(word string) {
	node := ac.root
	for _, ch := range word {
		if _, ok := node.children[ch]; !ok {
			//node.children[ch] = NewACNode(string(ch))
			node.children[ch] = NewACNode()
		}
		node = node.children[ch]
	}
	node.isEnd = true
	node.word = word
}

func (ac *ACAutomaton) Build() {
	queue := make([]*ACNode, 0)
	ac.root.fail = nil
	queue = append(queue, ac.root)
	for len(queue) > 0 {
		node := queue[0]
		queue = queue[1:]
		for ch, child := range node.children {
			if node == ac.root {
				child.fail = ac.root
			} else {
				failNode := node.fail
				for failNode != nil {
					if _, ok := failNode.children[ch]; ok {
						child.fail = failNode.children[ch]
						break
					}
					failNode = failNode.fail
				}
				if failNode == nil {
					child.fail = ac.root
				}
			}
			queue = append(queue, child)
		}
	}
}

func (ac *ACAutomaton) Search(text string) []string {
	result := make([]string, 0)
	node := ac.root
	for i, char := range text {
		for node.children[char] == nil && node != ac.root {
			node = node.fail
			if node.isEnd {
				result = append(result, node.word)
				//node = ac.root
				fmt.Println("now is ", i)
			}
		}
		node = node.children[char]
		if node == nil {
			node = ac.root
		}
		if node.isEnd {
			result = append(result, node.word)
			//node = ac.root
			fmt.Println("now is ", i)
		}
	}
	return result
}

func (ac *AC) Search(text string, replace rune) string {
	var result []rune
	node := ac.root
	for _, c := range text {
		for node != nil && !node.children[c] {
			node = node.fail
		}
		if node == nil {
			node = ac.root
		} else {
			node = node.children[c]
			if node.isEndOfWord {
				for i := len(result) - node.depth; i < len(result); i++ {
					result[i] = replace
				}
			}
		}
		result = append(result, c)
	}
	return string(result)
}

var (
	dirtyFiles = "simple/dirty_en.txt;simple/dirty_and_usa.txt;simple/dirty_korea.txt;simple/dirty4.txt;simple/dirty5.txt"
)

func main() {
	files := strings.Split(dirtyFiles, ";")
	ac := NewACAutomaton()
	for _, file := range files {
		h, err := os.Open(file)
		if err != nil {
			fmt.Println("open file: error:", file, err)
			continue
		}
		defer h.Close()

		scanner := bufio.NewScanner(h)
		for scanner.Scan() {
			line := scanner.Text()
			ac.AddWord(line)
		}
		if err := scanner.Err(); err != nil {
			fmt.Println("Error scanning file: ", err)
		}
	}
	ac.AddWord("中国")
	ac.AddWord("中国人")
	ac.AddWord("中国小说网")
	ac.Build()

	//fmt.Println("res:", ac.Search("我日卵子妈了个巴子杀人放火"))
	//fmt.Println("res:", ac.Search("我是中国人，中国小说网站有各类小说"))
	str := "我是中国人"
	str2 := "中国"
	fmt.Println([]rune(str), []rune(str2))
	fmt.Println("res:", ac.Search("我是中国人"))
	//fmt.Println("res:", ac.Search("中国人"))
}
